package com.yupi.yuaiagent.chatmemory;


import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.json.JSON;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.yupi.yuaiagent.domain.entity.AiChatMessage;
import com.yupi.yuaiagent.service.AiChatMessageService;
import dev.langchain4j.data.message.ChatMessage;
import jakarta.annotation.Resource;
import opennlp.tools.util.StringUtil;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * 基于MySQL持久化存储的对话记忆
 */
public class MySqlChatMemory implements ChatMemory {

    private final AiChatMessageService aiChatMessageService;

    public MySqlChatMemory(AiChatMessageService aiChatMessageService) {
        this.aiChatMessageService = aiChatMessageService;
    }


    /**
     * 存储到数据库
     * @param conversationId
     * @param messages
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        // id conversationId message createTime updateTime isDelete
        List<AiChatMessage> allMessage = messages.stream().map(message -> {
            // 把message转换成ChatMemory
            AiChatMessage chatMessage = new AiChatMessage();
            chatMessage.setConversationId(conversationId);
            chatMessage.setContent(message.getText());
            chatMessage.setType(message.getMessageType().toString());
            chatMessage.setCreateTime(LocalDateTime.now());
            chatMessage.setUpdateTime(LocalDateTime.now());
            chatMessage.setIsDelete(0);
            return chatMessage;
        }).collect(Collectors.toList());
        aiChatMessageService.saveBatch(allMessage);
    }

    /**
     * 获取当前会话id的最近n条消息 （可以配置）
     * @param conversationId
     * @param lastN
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<AiChatMessage> allMessages = aiChatMessageService.list(
                new QueryWrapper<AiChatMessage>().eq("conversation_id", conversationId).orderByDesc("update_time")
        );

        // 转换为Message对象
        List<Message> messages = new ArrayList<>();
        for (AiChatMessage chatMessage : allMessages) {
            String type = chatMessage.getType();
            switch (type) {
                case "USER":
                    messages.add(new UserMessage(chatMessage.getContent())) ;
                    break;
                case "ASSISTANT":
                    messages.add(new AssistantMessage(chatMessage.getContent())) ;
                    break;
                case "SYSTEM":
                    messages.add(new SystemMessage(chatMessage.getContent()));
                    break;
                default:
                    throw new IllegalArgumentException("Unknown message type: " + type);
            }
        }
        // 获取当前会话id的最近n条消息 （可以配置）
        messages = messages.stream().skip(Math.max(0, messages.size() - lastN)).toList();
        return messages;
    }

    @Override
    public void clear(String conversationId) {
        boolean isDel = aiChatMessageService.remove(new QueryWrapper<AiChatMessage>().eq("conversation_id", conversationId));
    }
}
